Skip to content

State checkpointer refactor#394

Merged
fatemetkl merged 24 commits into
mainfrom
ft/state-checkpointing-refactor
Jun 19, 2025
Merged

State checkpointer refactor#394
fatemetkl merged 24 commits into
mainfrom
ft/state-checkpointing-refactor

Conversation

@fatemetkl

@fatemetkl fatemetkl commented Jun 3, 2025

Copy link
Copy Markdown
Collaborator

PR Type

Fix

Short Description

Clickup Ticket(s): https://app.clickup.com/t/868d4w5rt

Unifying our PerRoundStateCheckpointer and the Snapshotting functionality that supports early stopping.
This PR adds checkpointing.state_checkpointer modules with classes that can be used to checkpoint the state of the server and clients.

Tests Added

  • tests/checkpointing/test_state_checkpointer.py
  • tests/utils/test_early_stopper.py

I also manually tested our nnunet_example to make sure NnUnetServerCheckpointAndStateModule works as expected. We have a smoke test (load_from_checkpoint_example) for ServerPerRoundStateCheckpointer.

@codecov

codecov Bot commented Jun 3, 2025

Copy link
Copy Markdown

Codecov Report

Attention: Patch coverage is 94.31818% with 15 lines in your changes missing coverage. Please review.

Project coverage is 77.63%. Comparing base (2444580) to head (24581fb).
Report is 25 commits behind head on main.

Files with missing lines Patch % Lines
fl4health/checkpointing/state_checkpointer.py 93.22% 12 Missing ⚠️
fl4health/checkpointing/client_module.py 90.00% 1 Missing ⚠️
fl4health/checkpointing/server_module.py 94.11% 1 Missing ⚠️
fl4health/servers/nnunet_server.py 66.66% 1 Missing ⚠️
Additional details and impacted files
@@            Coverage Diff             @@
##             main     #394      +/-   ##
==========================================
+ Coverage   76.83%   77.63%   +0.79%     
==========================================
  Files         156      157       +1     
  Lines        9472     9553      +81     
==========================================
+ Hits         7278     7416     +138     
+ Misses       2194     2137      -57     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

@fatemetkl fatemetkl requested a review from emersodb June 5, 2025 13:01
Comment thread tests/smoke_tests/run_smoke_test.py
self.state_checkpoint_name
) # self.state_checkpoint_name initialized in base FLServer Class
):
self._load_server_state()

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We don't need to call _load_server_state() here because it will be called in FlServer's fit_with_per_round_checkpointing()

@fatemetkl fatemetkl requested review from lotif and nerdai June 5, 2025 13:36
@fatemetkl fatemetkl force-pushed the ft/state-checkpointing-refactor branch 2 times, most recently from b5c6a6d to 80fbafd Compare June 10, 2025 15:32
@fatemetkl fatemetkl force-pushed the ft/state-checkpointing-refactor branch from 019bc21 to 96f1858 Compare June 10, 2025 17:02

@emersodb emersodb left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nearly ready to go I think a couple of small questions and suggested refactors, but looks great!

Comment thread examples/nnunet_example/client.py
Comment thread fl4health/checkpointing/state_checkpointer.py Outdated
Comment thread examples/nnunet_example/server.py
Comment thread fl4health/checkpointing/state_checkpointer.py Outdated
Comment thread fl4health/checkpointing/state_checkpointer.py Outdated
Comment thread fl4health/checkpointing/state_checkpointer.py Outdated
Comment thread fl4health/checkpointing/state_checkpointer.py Outdated
Comment thread fl4health/utils/early_stopper.py Outdated
Comment thread fl4health/utils/snapshotter.py Outdated
Comment thread tests/utils/snapshotter_test.py
@emersodb emersodb self-requested a review June 19, 2025 17:32

@emersodb emersodb left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good to go!

@fatemetkl fatemetkl merged commit c937a32 into main Jun 19, 2025
10 checks passed
@fatemetkl fatemetkl deleted the ft/state-checkpointing-refactor branch June 19, 2025 18:41
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants